#!/usr/bin/env ruby

module Panoptimon
  module Collector
    class IOStat

      attr_reader :iostat, :flags, :opt

      def initialize (options={})
        @opt = options
        @iostat = opt[:iostat] || 'iostat'
      end

      def cmd
        cmd = [iostat, flags]

        if opt[:interval]
          cmd.push(opt[:interval].to_s)
          cmd.push(opt[:count].to_s) if opt[:count] # just for debugging
        end

        return cmd
      end

      def parse_header (l)
        head = l.chomp.split(/\s+/)
        head = Hash[*head.zip(0..(head.length-1)).flatten]
        want.values.find_all{|x| not(head[x])}.tap {|missed|
          warn "missing headers: "+ missed.join(', ') if missed.length > 0
        }
        o = {}; want.each {|k,v| o[k] = head[v] if head[v]}
        return o
      end

      def run
        p = IO.popen(cmd, 'r')
        if p.eof?
          raise $?;
        end
        cue(p)

        puts '{}' # beep

        omap = nil
        device_idx = nil

        until p.eof?
          l = getlines(p)

          # TODO something less sprawled
          omap ||= parse_header(l[0])
          device_idx ||= omap.delete(:device)

          puts JSON::generate(Hash[*l.drop(1).map {|x|
            r = x.split(/\s+/)
            [r[device_idx],
              Hash[*omap.keys.map {|k| [k, r[omap[k]].to_f]}.flatten]]
          }.flatten])

        end
      end
      
      class OS_linux < Panoptimon::Collector::IOStat
        def flags; '-xdk' ; end

        def cue (p)
          2.times { x = p.readline until x == "\n" } # header + first sample
        end

        def getlines (p)
          p.readline("\n\n").split(/\n/)
        end

        def want
          { device: 'Device:',
            'kb_read/s'  => 'rkB/s',
            'kb_write/s' => 'wkB/s',
            'rrqm/s'     => 'rrqm/s',
            'wrqm/s'     => 'wrqm/s',
            'reads/s'    => 'r/s',
            'writes/s'   => 'w/s',
            'avgrq-sz'   => 'avgrq-sz',
            'avgqu-sz'   => 'avgqu-sz',
            'await'      => 'await',
            'r_await'    => 'r_await',
            'w_await'    => 'w_await',
            'util'       => '%util',
          }
        end

      end
      class OS_freebsd < Panoptimon::Collector::IOStat::OS_linux; end  
      class OS_solaris < Panoptimon::Collector::IOStat
        def flags; '-xne' ; end

        def cue (p)
          2.times { x = p.readline until x =~ /^\s*extended / }
        end

        def getlines (p)
          lines = []
          while line = p.gets
            return lines if line =~ /^\s* extended /
            lines << line.chomp
          end
          # eof?
          return lines
        end

        def want
          { device: 'device',
            'kb_read/s'  => 'kr/s',
            'kb_write/s' => 'kw/s',
            'reads/s'    => 'r/s',
            'writes/s'   => 'w/s',
            'wait'       => 'wait',
            'wsvc_t'     => 'wsvc_t',
            'asvc_t'     => 'asvc_t',
            'pct_wait'   => '%w',
            'pct_busy'   => '%b',
            'errors'     => 'tot',
          }
        end
      end
      class OS_openbsd < Panoptimon::Collector::IOStat
        # default output format, -w required
        # TODO reimplement cmd -- will need -w before count
        # TODO reimplement run -- parsing column-wise
      end
    end
  end
end

########################################################################
require 'panoptimon/util'

$stdout.sync = true # persistent process

require 'json'
opt = ARGV[0] ?
    JSON::parse(ARGV[0], {symbolize_names: true})
  : {interval: 1, count: 2}

os_class = Panoptimon::Collector::IOStat.const_get(
  'OS_' + Panoptimon::Util.os.to_s)
os_class.new(opt).run

